import sys
import torch
sys.path.append('..')
from sent2vec import Sent2VecEmbeddings
import json
import numpy as np
import os
import faiss
import time
# BAAI/bge-large-zh-v1.5
import csv

with open('/home/lxy/wxbdata/NiuLM_图书馆测试.csv', 'r') as file:
    # 创建CSV阅读器对象
    csv_reader = csv.reader(file)
    corpus=[row[0] for row in csv_reader]

corpus=corpus[1:]
print(corpus,isinstance(corpus[0],str))

print('测试集',len(corpus))

#导入向量数据库raw数据的地方
with open('/home/lxy/wxbdata/merge.json','r',encoding='utf-8') as f:
    data=json.load(f)
    data=[d['answer'] for d in data]
print('向量数据库生数据',len(data))



index_dir='/data/lxy/RAT/10m9d_wxb_bge1.5_hnswivf_20500'
maxNorm=torch.load(os.path.join(index_dir,'maxNorm.pt'))
index=faiss.read_index(os.path.join(index_dir,'index.faiss'))
assert index.ntotal==len(data),f'index.index.ntotal={index.ntotal},len(data)={len(data)}\n可能是用add方法往里面加数据了，这样可能导致索引和生数据没法对齐，想想咋搞'

print('要开始了')
model=Sent2VecEmbeddings(model_name='BAAI/bge-large-zh-v1.5')
print('模型加载完了')



query_embedding,_ = model.embed_documents(corpus)
query_embedding=query_embedding.to('cpu')/maxNorm
print(query_embedding.shape)


D,I=index.search(query_embedding,k=30)
# print(query_embedding)
# print(D,I)
output_list=[]
k=3
with open('output.json','w',encoding='utf-8') as o:
    for i in range(len(I)):
        valid_answer=set()
        line={'question':corpus[i],}
        cnt=0
        for j in range(len(I[i])):
            if cnt>=k:
                break
            if data[I[i][j]] not in valid_answer:
                line[f'answer{cnt}']=data[I[i][j]]
                valid_answer.add(data[I[i][j]])
                cnt+=1
        output_list.append(line)
    # print(output_list[0])
    # json.dumps(output_list,o,ensure_ascii=False,indent=4)


import pandas as pd
df = pd.DataFrame(output_list)

# 创建一个Excel Writer对象，将DataFrame写入Excel文件
with pd.ExcelWriter('output.xlsx', engine='openpyxl') as writer:
    df.to_excel(writer, sheet_name='Sheet1', index=False)

print("Excel文件已创建成功！")



